import os, sys
import numpy as np
import scipy as sp
import matplotlib.pylab as plt
sys.path.append('./python')

from dcds_mask import *
from getOptWeightsOneFrame import *
from getOptWeightsTwoFrames import *

res_dir = 'weights_2fs'

###################################################################################################
if __name__ == '__main__':
    
    file_num = 180
    mask = DCDS_Mask()

    mask37 = mask.gen_mask(37)
    n_ref=37
    n_sig=37
    std_old_37 = []
    std_new_37 = []
    fig = plt.figure(figsize=(16,10))
    for i in range(16):
        print('> processing ch: {}'.format(i+1))
        os_wave_dir = 'OS_'+str(i+1)
        wfile = res_dir + '/' + 'w_os_' + str(n_ref) + '_' + str(i+1)+'.txt'
        std_oldx, std_newx, std_old2x, std_new2x = getWeightsTwoFrames(os_wave_dir, file_num=file_num, mask=mask37, n_ref=n_ref, n_sig=n_sig, out_fname=wfile)
        std_old_37.append([std_oldx, std_old2x])
        std_new_37.append([std_newx, std_new2x])
        tmp = np.loadtxt(wfile)
        plt.subplot(4,4,i+1)
        plt.plot(tmp,'--.')
        plt.title('ch :'+str(i+1))
    
    plt.tight_layout()
    plt.savefig('w_'+str(n_ref)+'_2fs.png')


    mask41 = mask.gen_mask(41)
    n_ref=41
    n_sig=41
    std_old_41 = []
    std_new_41 = []
    fig = plt.figure(figsize=(16,10))
    for i in range(16):
        print('> processing ch: {}'.format(i+1))
        os_wave_dir = 'OS_'+str(i+1)
        wfile = res_dir + '/' + 'w_os_' + str(n_ref) + '_' + str(i+1)+'.txt'
        std_oldx, std_newx, std_old2x, std_new2x = getWeightsTwoFrames(os_wave_dir, file_num=file_num, mask=mask41, n_ref=n_ref, n_sig=n_sig, out_fname=wfile)
        std_old_41.append([std_oldx, std_old2x])
        std_new_41.append([std_newx, std_new2x])
        tmp = np.loadtxt(wfile)
        plt.subplot(4,4,i+1)
        plt.plot(tmp,'--.')
        plt.title('ch :'+str(i+1))
    
    plt.tight_layout()
    plt.savefig('w_'+str(n_ref)+'_2fs.png')
    
    mask45 = mask.gen_mask(45)
    n_ref=45
    n_sig=45
    std_old_45 = []
    std_new_45 = []
    fig = plt.figure(figsize=(16,10))
    for i in range(16):
        print('> processing ch: {}'.format(i+1))
        os_wave_dir = 'OS_'+str(i+1)
        wfile = res_dir + '/' + 'w_os_' + str(n_ref) + '_' + str(i+1)+'.txt'
        std_oldx, std_newx, std_old2x, std_new2x = getWeightsTwoFrames(os_wave_dir, file_num=file_num, mask=mask45, n_ref=n_ref, n_sig=n_sig, out_fname=wfile)
        std_old_45.append([std_oldx, std_old2x])
        std_new_45.append([std_newx, std_new2x])
        tmp = np.loadtxt(wfile)
        plt.subplot(4,4,i+1)
        plt.plot(tmp,'--.')
        plt.title('ch :'+str(i+1))
    
    plt.tight_layout()
    plt.savefig('w_'+str(n_ref)+'_2fs.png')

    plt.figure()
    plt.plot(std_old_37, 'r:o', label='std_old_37')
    plt.plot(std_old_41, 'g:d', label='std_old_41')
    plt.plot(std_old_45, 'b:s', label='std_old_45')

    plt.plot(std_new_37, 'r--o', label='std_new_37')
    plt.plot(std_new_41, 'g--d', label='std_new_41')
    plt.plot(std_new_45, 'b--s', label='std_new_45')
    
    # 未优化
    np.savetxt('std_old_37_2fs.txt0', std_old_37, fmt='%10.5f')
    np.savetxt('std_old_41_2fs.txt0', std_old_41, fmt='%10.5f')
    np.savetxt('std_old_45_2fs.txt0', std_old_45, fmt='%10.5f')
    
    # 优化后
    np.savetxt('std_new_37_2fs.txt0', std_new_37, fmt='%10.5f')
    np.savetxt('std_new_41_2fs.txt0', std_new_41, fmt='%10.5f')
    np.savetxt('std_new_45_2fs.txt0', std_new_45, fmt='%10.5f')
    
    plt.ylim([1.5,5])

    plt.legend()
    plt.savefig('RMS_2fs.png')
#    plt.show()
